#pragma once

#include <algorithm>
#include <eigen3/Eigen/Core>
#include <math.h>
#include <set>
#include <stdlib.h>
#include <string>
#include <unistd.h>
#include <vector>

namespace fusion_perception {

class KM {
public:
    KM(void) = default;
    ~KM(void) = default;

    template <typename T>
    bool get(const T& associationMat, std::vector<std::pair<int, int>>& matchCps, bool needReverse = false)
    {
        matchCps.clear();
        uSize_ = associationMat.rows();
        vSize_ = associationMat.cols();
        if (uSize_ > vSize_) {
            return false;
        }
        exU_ = new double[uSize_];
        exV_ = new double[vSize_];
        vMatched_ = new int[vSize_];
        std::fill(vMatched_, vMatched_ + vSize_, -1);
        memset(exV_, 0, vSize_ * sizeof(double));

        for (int i = 0; i < uSize_; ++i) {
            exU_[i] = associationMat(i, 0);
            for (int j = 1; j < vSize_; ++j) {
                exU_[i] = std::max((float)exU_[i], associationMat(i, j));
            }
        }

        for (int i = 0; i < uSize_; ++i) {
            if (exU_[i] <= 0) {
                if (needReverse) {
                    matchCps.push_back(std::make_pair(-1, i));
                } else {
                    matchCps.push_back(std::make_pair(i, -1));
                }
                continue;
            }
            vSlack_ = new double[vSize_];
            std::fill(vSlack_, vSlack_ + vSize_, 999999999999);

            while (1) {
                usedU_.clear();
                usedV_.clear();

                if (findCP(associationMat, i)) {
                    break;
                }

                double d = 999999999999;
                for (int j = 0; j < vSize_; ++j)
                    if (usedV_.find(j) == usedV_.end()) {
                        d = std::min(d, vSlack_[j]);
                    }
                for (auto it = usedU_.begin(); it != usedU_.end(); it++) {
                    exU_[*it] -= d;
                }

                for (int j = 0; j < vSize_; ++j) {
                    if (usedV_.find(j) != usedV_.end()) {
                        exV_[j] += d;
                    } else {
                        vSlack_[j] -= d;
                    }
                }
            }
            delete[] vSlack_;
        }

        if (needReverse) {
            for (int j = 0; j < vSize_; ++j) {
                if (vMatched_[j] == -1) {
                    matchCps.push_back(std::make_pair(j, -1));
                } else if (associationMat(vMatched_[j], j) > 0) {
                    matchCps.push_back(std::make_pair(j, vMatched_[j]));
                } else {
                    matchCps.push_back(std::make_pair(-1, vMatched_[j]));
                    matchCps.push_back(std::make_pair(j, -1));
                }
            }
        } else {
            for (int j = 0; j < vSize_; ++j) {
                if (vMatched_[j] == -1) {
                    matchCps.push_back(std::make_pair(-1, j));
                } else if (associationMat(vMatched_[j], j) > 0) {
                    matchCps.push_back(std::make_pair(vMatched_[j], j));
                } else {
                    matchCps.push_back(std::make_pair(vMatched_[j], -1));
                    matchCps.push_back(std::make_pair(-1, j));
                }
            }
        }

        delete[] exU_;
        delete[] exV_;
        delete[] vMatched_;
        return true;
    }

private:
    template <typename T>
    bool findCP(const T& mat, int i)
    {
        usedU_.insert(i);
        for (int j = 0; j < vSize_; ++j) {
            if (usedV_.find(j) != usedV_.end())
                continue;
            double gap = exU_[i] + exV_[j] - mat(i, j);
            if (gap <= 0) {
                usedV_.insert(j);
                bool match_succ = vMatched_[j] == -1 || findCP(mat, vMatched_[j]);
                if (match_succ) {
                    vMatched_[j] = i;
                    return true;
                }
            } else {
                vSlack_[j] = std::min(vSlack_[j], gap);
            }
        }
        return false;
    }

private:
    int uSize_;
    int vSize_;
    double* exU_;
    double* exV_;
    int* vMatched_;
    double* vSlack_;
    std::set<int> usedU_;
    std::set<int> usedV_;
};

}
